package sortedset

type Consumer func(*Elem) bool

type SortedSet struct {
	dict     map[string]*Elem
	skipList *skipList
}

func New() *SortedSet {
	return &SortedSet{
		dict:     make(map[string]*Elem),
		skipList: newSkipList(),
	}
}

func (s *SortedSet) Len() int64 {
	return int64(len(s.dict))
}

func (s *SortedSet) PutOrSet(key string, score float64) int64 {
	elem, exists := s.dict[key]
	s.dict[key] = &Elem{
		Key:   key,
		Score: score,
	}
	if exists {
		if score != elem.Score {
			s.skipList.remove(key, score)
			s.skipList.insert(key, score)
		}
		return 1
	}
	s.skipList.insert(key, score)
	return 0
}

func (s *SortedSet) Remove(key string) int64 {
	elem, exists := s.dict[key]
	if !exists {
		return 0
	}
	s.skipList.remove(key, elem.Score)
	delete(s.dict, key)
	return 1
}

func (s *SortedSet) GetRank(key string, desc bool) int64 {
	elem, exists := s.dict[key]
	if !exists {
		return -1
	}
	r := s.skipList.getRank(key, elem.Score)
	if desc {
		return s.skipList.length - r
	}
	return r - 1
}

func (s *SortedSet) ForEachByRank(start int64, stop int64, desc bool, consumer Consumer) bool {
	if start < 0 {
		start = 0
	}
	if stop > s.Len() {
		stop = s.Len()
	}
	if start >= s.Len() || stop < start {
		return false
	}

	node := s.skipList.header.forwards[0].node
	if start > 0 {
		node = s.skipList.getByRank(start + 1)
	}
	if desc {
		node = s.skipList.trailer
		if start > 0 {
			node = s.skipList.getByRank(s.Len() - start)
		}
	}

	for i := 0; i < int(stop-start); i++ {
		if !consumer(&node.Elem) {
			return false
		}
		if desc {
			node = node.backward
			continue
		}
		node = node.forwards[0].node
	}
	return true
}

func (s *SortedSet) RangeByRank(start int64, stop int64, desc bool) []*Elem {
	if start < 0 {
		start = 0
	}
	if stop > s.Len() {
		stop = s.Len()
	}
	if start >= s.Len() || stop < start {
		return nil
	}
	elems := make([]*Elem, 0, stop-start)
	s.ForEachByRank(start, stop, desc, func(elem *Elem) bool {
		elems = append(elems, elem)
		return true
	})
	return elems
}

func (s *SortedSet) Count(min ScoreBorder, max ScoreBorder) int64 {
	var count int64 = 0
	s.ForEachByRank(0, s.Len(), false, func(elem *Elem) bool {
		score := elem.Score
		if min.GT(score) || max.LT(score) {
			return false
		}
		count++
		return true
	})
	return count
}

func (s *SortedSet) ForEachByScore(min ScoreBorder, max ScoreBorder, offset int64, limit int64, desc bool, consumer Consumer) bool {
	node := s.skipList.getFirstInScoreRange(min, max)
	if desc {
		node = s.skipList.getLastInScoreRange(min, max)
	}
	for node != nil && offset > 0 {
		cur := node
		node = cur.forwards[0].node
		if desc {
			node = cur.backward
		}
		offset--
	}
	for i := int64(0); (i < limit || limit < 0) && node != nil; i++ {
		if !consumer(&node.Elem) {
			return false
		}
		cur := node
		node = cur.forwards[0].node
		if desc {
			node = cur.backward
		}
		if node == nil || min.LT(node.Elem.Score) || max.GT(node.Elem.Score) {
			return false
		}
	}
	return true
}

func (s *SortedSet) RangeByScore(min ScoreBorder, max ScoreBorder, offset int64, limit int64, desc bool) []*Elem {
	if limit == 0 || offset < 0 {
		return nil
	}
	elems := make([]*Elem, 0)
	s.ForEachByScore(min, max, offset, limit, desc, func(elem *Elem) bool {
		elems = append(elems, elem)
		return true
	})
	return elems
}

func (s *SortedSet) RemoveByScore(min ScoreBorder, max ScoreBorder) int64 {
	removedKeys := s.skipList.RemoveRangeByScore(min, max)
	for _, key := range removedKeys {
		delete(s.dict, key)
	}
	return int64(len(removedKeys))
}

func (s *SortedSet) RemoveByRank(start int64, stop int64) int64 {
	removedKeys := s.skipList.RemoveRangeByRank(start, stop)
	for _, key := range removedKeys {
		delete(s.dict, key)
	}
	return int64(len(removedKeys))
}
